feat(Model Support): add Krea-2-Turbo model + LoRA support (WIP)#9304
Draft
Pfannkuchensack wants to merge 12 commits into
Draft
feat(Model Support): add Krea-2-Turbo model + LoRA support (WIP)#9304Pfannkuchensack wants to merge 12 commits into
Pfannkuchensack wants to merge 12 commits into
Conversation
Integrate Krea-2-Turbo (krea/Krea-2-Turbo) text-to-image per
NEW_MODEL_INTEGRATION.md: Krea2Transformer2DModel (single-stream MMDiT)
+ Qwen3-VL text encoder (12-layer hidden-state tap, 4D prompt_embeds)
+ reused Qwen-Image VAE + FlowMatchEulerDiscrete scheduler.
Backend:
- taxonomy: BaseModelType.Krea2, ModelType/ModelFormat.Qwen3VLEncoder,
Krea2VariantType (Turbo = "krea2_turbo" to avoid Z-Image collision)
- config probes: Main_Diffusers/Checkpoint_Krea2, Qwen3VLEncoder,
LoRA_LyCORIS_Krea2 (text_fusion/time_mod_proj signature; excluded
from the Qwen-Image probe to avoid double-match)
- loaders for the diffusers pipeline + standalone Qwen3-VL encoder,
with runtime workarounds for the HF model's version mismatches
(AutoTokenizer, extra_special_tokens={}, rope_parameters->rope_scaling)
- native sampling (pack/unpack, position_ids, linear-mu shift) and
hand-written Euler denoise loop; reuses qwen_image l2i/i2l
- invocations: model_loader, text_encoder, denoise, lora_loader, plus
two ecosystem enhancers (conditioning rebalance, seed variance)
- LoRA conversion for diffusers PEFT (lora_transformer- prefix)
Frontend:
- 'krea-2' base + qwen3_vl_encoder type/format across model maps,
buildKrea2Graph, addKrea2LoRAs, graph-builder denoise/base lists,
optimal dimension 1024, regenerated schema.ts
Fixes:
- estimate transformer working memory in krea2_denoise so the cache
reserves activation headroom and offloads more model under partial
loading; fixes fp8 + LoRA OOM at 1024 (model was placed before LoRA
patches were applied, leaving no room for their activations)
WIP: requires diffusers main (>=0.39 dev) for Krea2Transformer2DModel;
pyproject.toml temporarily pins diffusers to git main.
github-actions
Bot
added
api
python
Collaborator
|
Amazing! I was just thinking of working on this myself and you did it for me! |
Allow non-diffusers Krea-2 transformers (GGUF/fp8) to run with standalone single-file VAE + Qwen3-VL encoder, fixing several blockers found in testing. - buildKrea2Graph: drop the hard "requires Diffusers-format" assert; instead require both a VAE and a Qwen3-VL encoder to be selected when the transformer is not diffusers (mirrors readiness.ts). - Qwen3-VL encoder remap: handle both single-file key conventions — implicit (model.layers.*) and explicit (model.language_model.*). The old blind model.* -> language_model.* turned the bf16 file's keys into language_model.language_model.* (398 meta tensors -> "Cannot copy out of meta tensor" crash). Both files now load 0 missing / 0 unexpected / 0 meta. - Qwen3-VL tokenizer/config: broaden the offline-cache fallback from OSError to Exception so a partial HF cache (config present, vocab missing) re-fetches instead of dying with TypeError. - Qwen3-VL encoder fp8: keep an fp8 source checkpoint fp8-resident with per-layer upcast (storage float8_e4m3fn, compute bf16) instead of dequantizing to bf16. Halves resident VRAM (~8.9GB -> ~4.4GB), avoiding partial-load thrashing alongside a large transformer. Auto-enabled for fp8 sources on CUDA; bf16 files stay bf16. - Qwen-Image VAE: a native-layout qwen_image_vae single file is classified with the Anima base and loaded as AutoencoderKLWan, but the qwen l2i/i2l nodes need AutoencoderKLQwenImage. Add backend/krea2/vae_compat.py::as_qwen_image_vae to reinterpret a Wan VAE as AutoencoderKLQwenImage (state dicts are identical, 194/194 keys); both qwen VAE nodes use it. Idempotent for real QwenImage VAEs.
Contributor
|
Can you add it to the list of supported models here https://github.com/invoke-ai/InvokeAI/blob/main/README.md?plain=1#L61 |
An upstream merge reintroduced the AutoencoderKLQwenImage isinstance asserts in the qwen VAE nodes (without the import → F821) and dropped the adapter in the i2l path. A native-layout qwen_image_vae single file is classified with the Anima base and loaded as AutoencoderKLWan, so the asserts fail at runtime. - qwen_image_latents_to_image: drop the reintroduced pre-device assert (the as_qwen_image_vae adapter inside model_on_device already handles the class). - qwen_image_image_to_latents: restore the as_qwen_image_vae import + adapter call, remove both asserts. - estimate_vae_working_memory_qwen_image only reads tensor shape + element size, so it runs correctly on either VAE class before the adapter.
…l_loader
krea2_model_loader was added to MainModelLoaderNodes but not to
isMainModelWithoutUnet, and the guard wasn't a type predicate — so it never
narrowed modelLoader. OutputFields of the loader union collapses to the common
'vae' field, making g.addEdge(modelLoader, 'unet', ...) in addInpaint/addOutpaint
fail to type-check ('unet' not assignable to 'vae').
Redefine the guard as a type predicate keyed on the inverse (only
main_model_loader/sdxl_model_loader expose a unet), so every transformer-based
loader is treated as unet-less automatically and the negated branch narrows to
the unet-bearing loaders.
zKrea2VariantType was only used within common.ts (in zAnyModelVariant) and never referenced externally, so knip flagged it as an unused export. Every sibling variant enum avoids this by being asserted in common.test-d.ts; add the missing Krea2VariantType assertion, which both uses the export and verifies the manual zod enum matches the generated S['Krea2VariantType'].
Pfannkuchensack
requested review from
JPPhoto,
blessedcoolant,
dunkeroni and
lstein
as code owners
diffusers 0.39.0 is the first stable release containing Krea2Pipeline / Krea2Transformer2DModel (plus the Qwen-Image VAE and Qwen3-VL text encoder Krea-2 relies on). Replace the temporary git-main dependency with the pinned release and update the lockfile's diffusers entry (version, sdist, wheel, specifier) to the official PyPI 0.39.0 artifacts.
Close the remaining gaps against docs/new-model-integration for Krea-2. Metadata recall (§7): buildKrea2Graph now records the standalone VAE, Qwen3-VL encoder, and both conditioning enhancers (seed-variance + rebalance) to image metadata, and parsing.tsx adds the matching recall handlers (base-guarded to 'krea-2'), so a Krea-2 image's VAE/encoder/enhancer settings restore on recall — important for reproducing single-file/GGUF generations. Starter models: add Krea-2 Raw (Base variant, full pipeline), Krea-2 Turbo GGUF Q4_K_M / Q8_0 (vantagewithai/Krea-2-Turbo-GGUF) with Qwen-Image VAE + Qwen3-VL encoder dependencies, and a standalone Qwen3-VL 4B encoder (Qwen/Qwen3-VL-4B- Instruct). The VAE dependency reuses the existing diffusers qwen_image_vae starter; krea2_turbo gains its explicit Turbo variant. Tests: add config-probe unit tests for Krea-2 variant detection (name heuristic, _has_krea2_keys, GGUF/checkpoint/diffusers variant, default settings) and for the single-file Qwen3-VL encoder probe (visual-tower vs. text-only Qwen3). 31 tests.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Integrate Krea-2 text-to-image per
NEW_MODEL_INTEGRATION.md: both Krea-2-Turbo (krea/Krea-2-Turbo, distilled) and Krea-2-Raw (krea/Krea-2-Raw, undistilled Base). Architecture:Krea2Transformer2DModel(single-stream MMDiT, ~12B) + Qwen3-VL text encoder (12-layer hidden-state tap → 4Dprompt_embeds) + reused Qwen-Image VAE (AutoencoderKLQwenImage) +FlowMatchEulerDiscreteScheduler.is_distilled=true): fixedmu=1.15, 8 steps, CFG off (cfg 1.0).is_distilled=false): resolution-awaremu, ~28 steps, CFG ~4.5. Variant is read from the pipelineis_distilledflag (single-file/GGUF fall back to a filename heuristic).Formats: full Diffusers pipeline, single-file checkpoint (incl. ComfyUI scaled fp8), and GGUF (Q2–Q8). Single-file/GGUF ship only the transformer, so a standalone Qwen-Image VAE + Qwen3-VL encoder are selected in the loader UI (enforced by readiness before enqueue). NVFP4 intentionally skipped (needs Blackwell FP4 kernels).
VRAM: fp8 layerwise-cast weight storage for the transformer (diffusers + single-file) and for the fp8 Qwen3-VL encoder (~8.9 GB bf16 → ~4.4 GB resident), keeping 1024² + LoRA within 24 GB.
diffusers dependency (resolved)
Krea2Transformer2DModel/Krea2Pipelinelanded in diffusers 0.39.0 (stable).pyproject.tomlnow pinsdiffusers[torch]==0.39.0anduv.lockis updated accordingly — the previous git-main blocker is gone.Notable extras
isMainModelWithoutUnetis now a proper type predicate covering all transformer-based loaders).QA Instructions
uv sync --extra cuda(pulls diffusers 0.39.0); confirmpython -c "from diffusers import Krea2Transformer2DModel, Krea2Pipeline".Krea-2-Turboand aKrea-2-Rawfolder. Confirm they probe asmain / diffusers / krea-2with variantkrea2_turbo/krea2_base, and the bundled encoder asqwen3_vl_encoder. On model select, Turbo defaults to 8 steps / cfg 1.0, Raw to ~28 steps / cfg 4.5.vantagewithai/Krea-2-Turbo-GGUF) + the standalone Qwen-Image VAE + Qwen3-VL encoder. Confirm the loader UI requires the VAE + encoder before enqueue and that generation succeeds. The fp8 encoder logsFP8 layerwise casting enabled for Qwen3-VL encoder.lora.lycoris.krea-2(not qwen-image), applies, and that 1024² + LoRA + fp8 does not OOM.Tests
_has_krea2_keys, GGUF/checkpoint/diffusers variant, default settings) and the single-file Qwen3-VL encoder probe (visual-tower vs. text-only Qwen3). Run:pytest tests/backend/model_manager/configs/test_krea2_main_config.py tests/backend/model_manager/configs/test_qwen3_vl_encoder_config.py.Merge Plan
diffusers blocker is resolved (pinned to stable 0.39.0). No DB schema changes.
paramsSlicegains Krea-2 fields with a corresponding migration.Out of scope / follow-ups
Checklist
What's Newcopy (if doing a release after this PR)